import torch
import torch.nn as nn
from torchvision import models


class WideResNet50(nn.Module):
    def __init__(self, class_name, use_pretrained=True):
        super(WideResNet50, self).__init__()
        resnet50 = models.wide_resnet50_2(pretrained=True)
        fc_in = resnet50.fc.in_features
        fc_out = 7

        for idx, p in enumerate(resnet50.parameters()):
            if idx < 96:
                p.requires_grad = False

        resnet50.fc = nn.Sequential(
            nn.Linear(fc_in, 256),
            nn.SELU(),
            nn.Linear(256, fc_out))

        if use_pretrained:
            resnet50.load_state_dict(torch.load(
                "../saved/resnet50-flora-classi.pth"))

        # Build extractor
        self.lte = nn.Sequential()
        self.lte.add_module("conv1", resnet50.conv1)
        self.lte.add_module("bn1", resnet50.bn1)
        self.lte.add_module("relu", resnet50.relu)
        self.lte.add_module("maxpool", resnet50.maxpool)
        self.lte.add_module("layer1", nn.Sequential())
        self.lte.layer1.add_module("0", resnet50.layer1[0])
        self.lte.layer1.add_module("1", resnet50.layer1[1])
        self.lte.layer1.add_module("2", resnet50.layer1[2])
        # self.lte.layer1.add_module("2", torch.utils.bottleneck())
        # self.lte.layer1[2].add_module("conv1", resnet50.resnet50.layer1[2].conv1)
        # self.lte.layer1[2].add_module("bn1", resnet50.resnet50.layer1[2].bn1)
        # self.lte.layer1[2].add_module("conv1", resnet50.resnet50.layer1[2].conv1)

    def forward(self, x):
        return self.lte(x)
